package com.pai4j.aigc.llm;

import com.pai4j.aigc.config.properties.KimiLLMPropertiesConfig;
import com.pai4j.common.util.llm.MoonShotAIClient;
import com.pai4j.domain.vo.llm.ChatCompletionMessage;
import com.pai4j.domain.vo.llm.ChatCompletionRequest;
import com.pai4j.domain.vo.llm.ChatCompletionResponse;
import com.pai4j.domain.vo.llm.ChatCompletionStreamChoice;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.context.annotation.Primary;
import org.springframework.stereotype.Service;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import java.util.List;

/**
 *    kimi-ai
 */
@Service
@Primary
public class KimiAIService extends AbstractLLMChatService {

    @Autowired
    private KimiLLMPropertiesConfig kimiLLMPropertiesConfig;

    /**
     * 同步对话
     *
     * @param messageList
     * @return
     * @throws Exception
     */
    public ChatCompletionResponse chat(List<ChatCompletionMessage> messageList) throws Exception {
        MoonShotAIClient client = new MoonShotAIClient(kimiLLMPropertiesConfig.getApiKey());
        final List<ChatCompletionMessage> messages = messageList;
        return client.ChatCompletion(new ChatCompletionRequest(
                kimiLLMPropertiesConfig.getModel(),
                messages,
                4096,
                0.3f,
                1
        ));
    }

    @Override
    protected ChatCompletionResponse chat(List<ChatCompletionMessage> messageList, Object config) throws Exception {
        if (config instanceof LLMRuntimeConfig runtime && runtime.getApiKey() != null) {
            String apiKey = runtime.getApiKey();
            String model = runtime.getModel() != null ? runtime.getModel() : kimiLLMPropertiesConfig.getModel();
            MoonShotAIClient client = new MoonShotAIClient(apiKey);
            final List<ChatCompletionMessage> messages = messageList;
            Integer maxTokens = runtime.getMaxTokens() != null ? runtime.getMaxTokens() : 4096;
            Float temperature = runtime.getTemperature() != null ? runtime.getTemperature() : 0.3f;
            return client.ChatCompletion(new ChatCompletionRequest(
                    model,
                    messages,
                    maxTokens,
                    temperature,
                    1
            ));
        }
        return super.chat(messageList, config);
    }

    /**
     * 流式对话
     *
     * @param emitter
     * @param messageList
     * @throws Exception
     */
    public String chat(SseEmitter emitter, List<ChatCompletionMessage> messageList) throws Exception {

        StringBuilder aiMessage = new StringBuilder();

        MoonShotAIClient client = new MoonShotAIClient(kimiLLMPropertiesConfig.getApiKey());


        final List<ChatCompletionMessage> messages = messageList;
        try {
            client.ChatCompletionStream(new ChatCompletionRequest(
                    kimiLLMPropertiesConfig.getModel(),
                    messages,
                    4096,
                    0.5f,
                    1
            )).subscribe(
                    streamResponse -> {
                        if (streamResponse.getChoices().isEmpty()) {
                            return;
                        }
                        for (ChatCompletionStreamChoice choice : streamResponse.getChoices()) {
                            String finishReason = choice.getFinishReason();
                            if (finishReason != null) {
                                // capture usage if provider supplies it at finish
                                if (choice.getUsage() != null) {
                                    this.lastUsage = choice.getUsage();
                                }
                                emitter.send("finished");
                                continue;
                            }
                            String responseContent = choice.getDelta().getContent();
                            if (responseContent != null) {
                                responseContent = responseContent.replaceAll("\n", "<br>");
                                System.out.println(responseContent);
                                emitter.send(responseContent);
                                aiMessage.append(responseContent);
                            }
                            // Some providers send usage in an interim chunk
                            if (choice.getUsage() != null) {
                                this.lastUsage = choice.getUsage();
                            }
                        }
                    },
                    error -> {
                        error.printStackTrace();
                    },
                    () -> {
//                        emitter.complete(); // 完成事件流发送
                    }
            );
        } catch (Exception e) {
            e.printStackTrace();
        }
        return aiMessage.toString();
    }

    @Override
    protected String chat(SseEmitter emitter, List<ChatCompletionMessage> messageList, Object config) throws Exception {
        if (config instanceof LLMRuntimeConfig runtime && runtime.getApiKey() != null) {
            String apiKey = runtime.getApiKey();
            String model = runtime.getModel() != null ? runtime.getModel() : kimiLLMPropertiesConfig.getModel();
            StringBuilder aiMessage = new StringBuilder();
            MoonShotAIClient client = new MoonShotAIClient(apiKey);
            final List<ChatCompletionMessage> messages = messageList;
            try {
                Integer maxTokens = runtime.getMaxTokens() != null ? runtime.getMaxTokens() : 4096;
                Float temperature = runtime.getTemperature() != null ? runtime.getTemperature() : 0.5f;
                client.ChatCompletionStream(new ChatCompletionRequest(
                        model,
                        messages,
                        maxTokens,
                        temperature,
                        1
                )).subscribe(
                        streamResponse -> {
                            if (streamResponse.getChoices().isEmpty()) {
                                return;
                            }
                            for (ChatCompletionStreamChoice choice : streamResponse.getChoices()) {
                                String finishReason = choice.getFinishReason();
                                if (finishReason != null) {
                                    if (choice.getUsage() != null) {
                                        this.lastUsage = choice.getUsage();
                                    }
                                    emitter.send("finished");
                                    continue;
                                }
                                String responseContent = choice.getDelta().getContent();
                                if (responseContent != null) {
                                    responseContent = responseContent.replaceAll("\n", "<br>");
                                    emitter.send(responseContent);
                                    aiMessage.append(responseContent);
                                }
                                if (choice.getUsage() != null) {
                                    this.lastUsage = choice.getUsage();
                                }
                            }
                        },
                        error -> {
                            error.printStackTrace();
                        },
                        () -> {
                        }
                );
            } catch (Exception e) {
                e.printStackTrace();
            }
            return aiMessage.toString();
        }
        return super.chat(emitter, messageList, config);
    }
}